import numpy as np
import pylab as pl
import cv2.cv as cv # openCV
import numpy.random as nr
import os
import scipy
import math
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d as ip
from matplotlib.pyplot import savefig
import csv
from time import  time as tm
import time
import Image
import snake_energy
import copy

###### Function which converts cartesian coordinates into polar coordinates
def toPolar(x,y):
    r = math.sqrt(pow(x,2) + pow(y,2))
    if x == 0:
        x = 0.0001
    if(y < 0 and x < 0):
        theta = np.arctan2(abs(y),abs(x)) + np.pi
    elif(y < 0 and x > 0):
        theta = np.arctan2(abs(y),-x) + np.pi
    else:
        theta = np.arctan2(y,x)
    return (r,theta)


####### Function which converts polar coordinates into cartesian coordinates
def toCartesian(r, theta,(Mx,My)):
    x = (int)(r * np.cos(theta) + Mx)
    y = (int)(r * np.sin(theta) + My)
    return (x,y)

###### Function for initialization of the snake points
###### M = center; r = radius; N = number of points
def initPoints((Mx,My), r, N, snake):   
    for n in range(N):
        theta = n * (2*np.pi/N)
        snake.append(toCartesian(r,theta,(Mx,My)))
    return snake


def sort(polarSnake):
    sortPS = []
    preSort = []
    thetaVal = []
    for j in range(len(polarSnake)):
        thetaVal.append(polarSnake[j][1])
        
    ###### Rotation of Snake --> start an Minimum
    x = thetaVal.index(min(thetaVal))
    for i in (range(x,len(polarSnake))):
        preSort.append(polarSnake[i])
    for i in range(x):
        preSort.append(polarSnake[i])   

    ###### Twists in the Snake must be removed
    for i in range(len(preSort)):
        if i > 0:
            if(preSort[i-1][1] < preSort[i][1]) and (preSort[i][1] > sortPS[len(sortPS)-1][1]) and preSort[i][1]-sortPS[len(sortPS)-1][1] >= 0.01 and preSort[i][1]-sortPS[len(sortPS)-1][1] <= 1.5:   
                sortPS.append(preSort[i]) 
        else:
            sortPS.append(preSort[i])
    return sortPS

###### Function for interpolation between the snake points
def interpolate(snake,j):   
    splines = []
    Mx = 0
    My = 0
    for (x,y) in snake:
        Mx = Mx + x
        My = My + y
    Mx = Mx/len(snake)
    My = My/len(snake)
    
    shiftSnake = []
    for (x,y) in snake:
        shiftSnake.append((x-Mx,y-My))
    
    polarSnake = []
    for j in range(len(shiftSnake)):
        polarSnake.append(toPolar(shiftSnake[j][0],shiftSnake[j][1]))
    polarSnake = sort(polarSnake)
    polarSnake.append((polarSnake[0][0],(2*np.pi+polarSnake[0][1])))  

    rVal = []
    thetaVal = []
    for j in range(len(polarSnake)):
        rVal.append(polarSnake[j][0])
        thetaVal.append(polarSnake[j][1])

    f = ip(thetaVal, rVal, kind = 'cubic',bounds_error=False)    
#    xnew = np.linspace(0,7, 1000)
#    plt.plot(thetaVal,rVal,'o-',xnew,f(xnew),'-')
#    plt.legend(['data', 'cubic'], loc='best')
#    plt.show()
#    plt.savefig('plot.png')

    polarSpline = []

    for i in range(len(thetaVal)-1):
        a = np.int(round((np.rad2deg(thetaVal[i+1])-np.rad2deg(thetaVal[i]))))
        if a == 0:
            b = 0
        else:
            b = ((thetaVal[i+1]-thetaVal[i])/(a)) 

        newTheta = 0
        for j in range(a-1):
            newTheta = thetaVal[i] + (j*b)
            if newTheta <= thetaVal[0]:
                newTheta += 2*np.pi
                newR = (float)(f(newTheta))
                polarSpline.append((newR, newTheta-2*np.pi))
            else:
                newR = (float)(f(newTheta))
                polarSpline.append((newR, newTheta))
        
    for (r,theta) in polarSpline:
        splines.append(toCartesian(r, theta,(Mx,My)))
    return splines


####### Load image
image = cv.LoadImage('cirk.png')

######## general things
dst = cv.CreateImage(cv.GetSize(image), cv.IPL_DEPTH_8U, 1)
cv.CvtColor(image, dst, cv.CV_RGB2GRAY)
snake=[]
(w,h)=cv.GetSize(image)

####### Place initial points
initPoints((250,250),200,100,snake) 

####### Energy functionals, usually denoted alpha, beta and gamma.
a=[0.8]
b=[0.5]
c=[0.5]    

####### Mask for Snake
mask=cv.CloneImage(dst)

####### Loop the contour advacement
energy = np.zeros(100)
    

####### Loop the contour advacement
i=0
while i<100:
    shake = []
    frame = cv.CloneImage(dst)
    for (x,y) in snake:
     #######  add noice to the particle positions to escape local minima.
        shake.append(tuple((x,y)+nr.randint(-2,3,size=2))) 
        cv.Circle(frame, (x,y),2,200,cv.CV_FILLED,1)
                        
    #######  Interpolation between snake points  
    splines = interpolate(snake,i)
    for (x,y) in splines:
        cv.Circle(frame, (x,y),1,150,cv.CV_FILLED,1)    
            
    snake = shake
    
    ####### Save the images
    if i < 10:
        cv.SaveImage('frame00%s'%i+'.png',frame)
    elif i < 100:
        cv.SaveImage('frame0%s'%i+'.png',frame)
    else:
        cv.SaveImage('frame%s'%i+'.png',frame)

    ####### Snake energy
    E = snake_energy.snake_energy(snake,dst,(a[0],b[0],c[0]))
#   print E.E
    energy[i] = copy.copy(E.E)                     
        
    ####### Contour advancement
    snake = cv.SnakeImage(mask,snake,a,b,c,(15,15),(cv.CV_TERMCRIT_ITER,1, 0.01))
    i += 1
####### end

####### Plot energy
pl.plot(energy)
pl.savefig('energy.png')

####### create video from pictures
os.system('mencoder mf://*frame*.png -mf w=200:h=200:fps=4:type=png -ovc lavc \-lavcopts vcodec=mpeg4:mbd=2:trell -oac copy -o output.avi')
